"""
Event system integration with PyTorch Lightning callbacks.

This module provides a bridge between PyTorch Lightning's callback system
and our enhanced event-based architecture, translating Lightning events into our event types.
"""

from typing import Any, Dict, List, Optional, Union
import pytorch_lightning as pl
from torch.utils.data import DataLoader

from .core import (
    EventDispatcher,
    global_dispatcher,
    EventType,
    EventFactory,
    TaskEvent,
    TrainingEvent,
    EvaluationEvent,
)


class EventSystemCallback(pl.Callback):
    """
    PyTorch Lightning callback that dispatches events to our enhanced event system.

    This callback bridges the gap between PyTorch Lightning's callback system
    and our custom event system, translating Lightning callbacks into our event types.
    """

    def __init__(self, dispatcher: Optional[EventDispatcher] = None):
        """
        Initialize the event system callback.

        Args:
            dispatcher: Optional custom event dispatcher to use
                        If None, uses the global dispatcher
        """
        super().__init__()
        self.dispatcher = dispatcher or global_dispatcher
        self.current_task = 0

    def setup(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        stage: Optional[str] = None,
    ) -> None:
        """Called when fit, validate, test, predict begins."""
        event = EventFactory.create_training_event(
            event_type=EventType.TRAINING_START,
            pl_module=pl_module,
            trainer=trainer,
            source="EventSystemCallback.setup"
        )
        self.dispatcher.dispatch(event)
            TrainingEvent(
                event_type=EventType.TRAINING_START,
                pl_module=pl_module,
                trainer=trainer,
                source="EventSystemCallback",
            )
        )

    def teardown(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        stage: Optional[str] = None,
    ) -> None:
        """Called when fit, validate, test, predict ends."""
        self.dispatcher.dispatch(
            TrainingEvent(
                event_type=EventType.TRAINING_END,
                pl_module=pl_module,
                trainer=trainer,
                source="EventSystemCallback",
            )
        )

    def on_train_epoch_start(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
    ) -> None:
        """Called when the train epoch begins."""
        self.dispatcher.dispatch(
            TrainingEvent(
                event_type=EventType.EPOCH_START,
                pl_module=pl_module,
                trainer=trainer,
                current_epoch=trainer.current_epoch,
                task_id=self.current_task,
                source="EventSystemCallback",
            )
        )

    def on_train_epoch_end(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
    ) -> None:
        """Called when the train epoch ends."""
        self.dispatcher.dispatch(
            TrainingEvent(
                event_type=EventType.EPOCH_END,
                pl_module=pl_module,
                trainer=trainer,
                current_epoch=trainer.current_epoch,
                task_id=self.current_task,
                source="EventSystemCallback",
            )
        )

    def on_train_batch_start(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        batch: Any,
        batch_idx: int,
        unused: Optional[int] = 0,
    ) -> None:
        """Called when the train batch begins."""
        # Dispatch both batch available and batch start events
        self.dispatcher.dispatch(
            DataEvent(
                event_type=EventType.BATCH_AVAILABLE,
                batch=batch,
                batch_idx=batch_idx,
                source="EventSystemCallback",
            )
        )

        self.dispatcher.dispatch(
            TrainingEvent(
                event_type=EventType.BATCH_START,
                pl_module=pl_module,
                trainer=trainer,
                batch=batch,
                batch_idx=batch_idx,
                task_id=self.current_task,
                source="EventSystemCallback",
            )
        )

    def on_train_batch_end(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        outputs: Any,
        batch: Any,
        batch_idx: int,
        unused: Optional[int] = 0,
    ) -> None:
        """Called when the train batch ends."""
        self.dispatcher.dispatch(
            TrainingEvent(
                event_type=EventType.BATCH_END,
                pl_module=pl_module,
                trainer=trainer,
                batch=batch,
                outputs=outputs,
                batch_idx=batch_idx,
                task_id=self.current_task,
                source="EventSystemCallback",
            )
        )

    def on_validation_start(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
    ) -> None:
        """Called when validation begins."""
        self.dispatcher.dispatch(
            EvaluationEvent(
                event_type=EventType.EVALUATION_START,
                pl_module=pl_module,
                trainer=trainer,
                task_id=self.current_task,
                source="EventSystemCallback",
            )
        )

    def on_validation_end(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
    ) -> None:
        """Called when validation ends."""
        self.dispatcher.dispatch(
            EvaluationEvent(
                event_type=EventType.EVALUATION_END,
                pl_module=pl_module,
                trainer=trainer,
                task_id=self.current_task,
                source="EventSystemCallback",
            )
        )

    def on_test_start(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
    ) -> None:
        """Called when test begins."""
        self.dispatcher.dispatch(
            EvaluationEvent(
                event_type=EventType.EVALUATION_START,
                pl_module=pl_module,
                trainer=trainer,
                task_id=self.current_task,
                source="EventSystemCallback",
            )
        )

    def on_test_end(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule"
    ) -> None:
        """Called when test ends."""
        self.dispatcher.dispatch(
            EvaluationEvent(
                event_type=EventType.EVALUATION_END,
                pl_module=pl_module,
                trainer=trainer,
                task_id=self.current_task,
                source="EventSystemCallback",
            )
        )

    def on_task_start(
        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", task_id: int
    ) -> None:
        """
        Called when a new task begins.

        This is a custom method not part of PyTorch Lightning's callback system.
        It's expected to be called by the ContinualModule or another component.
        """
        self.current_task = task_id
        self.dispatcher.dispatch(
            TaskEvent(
                event_type=EventType.TASK_START,
                task_id=task_id,
                pl_module=pl_module,
                trainer=trainer,
                source="EventSystemCallback",
            )
        )

    def on_task_end(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        task_id: int,
        dataloader: Optional[DataLoader] = None,
    ) -> None:
        """
        Called when a task is completed.

        This is a custom method not part of PyTorch Lightning's callback system.
        It's expected to be called by the ContinualModule or another component.
        """
        self.dispatcher.dispatch(
            TaskEvent(
                event_type=EventType.TASK_END,
                task_id=task_id,
                pl_module=pl_module,
                trainer=trainer,
                dataloader=dataloader,
                source="EventSystemCallback",
            )
        )

    def on_metrics_computed(
        self,
        trainer: "pl.Trainer",
        pl_module: "pl.LightningModule",
        metrics: Dict[str, Any],
        task_id: Optional[int] = None,
    ) -> None:
        """
        Called when metrics are computed.

        This is a custom method not part of PyTorch Lightning's callback system.
        It's expected to be called by evaluation components.
        """
        self.dispatcher.dispatch(
            EvaluationEvent(
                event_type=EventType.METRICS_COMPUTED,
                metrics=metrics,
                pl_module=pl_module,
                trainer=trainer,
                task_id=task_id if task_id is not None else self.current_task,
                source="EventSystemCallback",
            )
        )


# Register this callback with the registry
from ..registry import CALLBACK_REGISTRY


@CALLBACK_REGISTRY.register(version="1.0.0")
class RegisteredEventSystemCallback(EventSystemCallback):
    """
    Registrable version of the EventSystemCallback.

    This allows the callback to be discovered and instantiated through the registry system.
    """

    @property
    def component_name(self) -> str:
        return "EventSystemCallback"

    @property
    def component_version(self) -> str:
        return "1.0.0"

    def get_config(self) -> Dict[str, Any]:
        return {}
